# 导入PyTorch核心库：用于构建神经网络、张量计算和自动求导
import torch
# 从torchvision导入数据集、预训练模型和数据预处理工具：
# datasets：提供经典数据集（如MNIST）；models：预训练模型（当前未用）；transforms：数据预处理函数
from torchvision import datasets, models, transforms
# 从torch.utils.data导入DataLoader：用于批量加载数据，支持多进程加速和打乱
from torch.utils.data import DataLoader
# 导入PyTorch的函数式API：提供激活函数（如ReLU）、池化等无状态操作
import torch.nn.functional as F
# 导入PyTorch优化器库：提供SGD、Adam等优化算法
import torch.optim as optim
# 导入matplotlib绘图库：用于绘制训练损失和准确率曲线
import matplotlib.pyplot as plt

# 设置中文显示：解决matplotlib绘图时中文乱码问题
# 1. 指定中文字体（兼容不同系统：SimHei（Windows）、WenQuanYi（Linux）、Heiti TC（macOS））
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
# 2. 解决负号显示为方块的问题
plt.rcParams["axes.unicode_minus"] = False

# -------------------------- 1. 数据准备：加载并预处理MNIST数据集 --------------------------
# 定义批次大小：每次训练/测试时加载64个样本，平衡训练速度和内存占用
batch_size = 64
# 定义数据预处理流水线：用transforms.Compose组合多个操作，确保训练/测试数据格式一致
transform = transforms.Compose([
    transforms.ToTensor(),  # 步骤1：将PIL图像（MNIST原始格式）转为PyTorch张量
    # 作用：1. 维度从(H,W)转为(C,H,W)（C=1，单通道灰度图）；2. 像素值从[0,255]归一化到[0,1]
    transforms.Normalize((0.1307,), (0.3081,))  # 步骤2：标准化（减均值、除标准差）
    # 注：(0.1307,)和(0.3081,)是MNIST数据集的全局均值和标准差，为官方统计值，用于消除数据分布差异，加速模型收敛
])

# 加载MNIST训练集：
train_dataset = datasets.MNIST(
    root='../dataset/mnist',  # 数据集保存路径（../表示上级目录，避免与代码文件混放）
    train=True,  # 指定为训练集（包含60000个样本）
    transform=transform,  # 应用上述预处理流水线
    download=True  # 若root路径下无数据集，自动从官网下载
)

# 构建训练集DataLoader：批量加载训练数据
train_loader = DataLoader(
    dataset=train_dataset,  # 传入训练集
    batch_size=batch_size,  # 每批次64个样本
    shuffle=True  # 训练时打乱样本顺序：避免模型学习数据顺序规律，提升泛化能力
)

# 加载MNIST测试集：（逻辑与训练集类似，区别在于train=False）
test_dataset = datasets.MNIST(
    root='../dataset/mnist',  # 与训练集同路径，避免重复下载
    train=False,  # 指定为测试集（包含10000个样本）
    download=True,  # 若已下载则自动跳过
    transform=transform  # 必须与训练集用相同预处理，保证数据格式一致
)

# 构建测试集DataLoader：
test_loader = DataLoader(
    dataset=test_dataset,
    batch_size=batch_size,
    shuffle=False  # 测试时无需打乱：只需按顺序评估所有样本即可
)


# -------------------------- 2. 模型构建：定义5层全连接神经网络 --------------------------
class Net(torch.nn.Module):  # 继承PyTorch的Module类（所有神经网络模型的基类）
    def __init__(self):  # 构造函数：初始化模型层结构
        super(Net, self).__init__()  # 调用父类构造函数，确保继承参数管理等核心功能

        # 定义全连接层（Linear层）：输入特征数→输出特征数，自动初始化权重w和偏置b
        self.fc1 = torch.nn.Linear(784, 512)  # 第1层：输入784（28×28像素展平后），输出512（压缩维度）
        self.fc2 = torch.nn.Linear(512, 256)  # 第2层：输入512（承接上一层输出），输出256
        self.fc3 = torch.nn.Linear(256, 128)  # 第3层：输入256，输出128
        self.fc4 = torch.nn.Linear(128, 64)  # 第4层：输入128，输出64
        self.fc5 = torch.nn.Linear(64, 10)  # 第5层（输出层）：输入64，输出10（对应10个数字类别0-9）

    def forward(self, x):  # 前向传播函数：定义数据在模型中的计算流程（核心）
        x = x.view(-1, 784)  # 步骤1：展平图像张量
        # 输入x形状为(batch_size, 1, 28, 28)，view(-1,784)将其转为(batch_size, 784)
        # -1表示自动计算该维度大小（保证总元素数不变），适配任意批次大小

        # 步骤2：全连接层+ReLU激活函数（增加非线性，让模型拟合复杂关系）
        x = F.relu(self.fc1(x))  # 第1层：fc1线性变换 → ReLU激活（将负数置0，保留正数）
        x = F.relu(self.fc2(x))  # 第2层：fc2 → ReLU
        x = F.relu(self.fc3(x))  # 第3层：fc3 → ReLU
        x = F.relu(self.fc4(x))  # 第4层：fc4 → ReLU

        # 步骤3：输出层（无激活）：直接返回fc5的线性输出（logits）
        # 原因：CrossEntropyLoss损失函数会自动对logits做Softmax，无需手动添加
        return self.fc5(x)


# 创建模型实例：自动初始化所有全连接层的权重（随机值）和偏置（默认0）
model = Net()

# -------------------------- 3. 损失函数与优化器：定义模型训练的“目标”和“更新规则” --------------------------
# 定义损失函数：交叉熵损失（CrossEntropyLoss），专为多分类任务设计
# 功能：1. 对模型输出的logits做Softmax（转为概率）；2. 计算与真实标签的交叉熵（衡量预测与真实的差距）
criterion = torch.nn.CrossEntropyLoss()

# 定义优化器：带动量的随机梯度下降（SGD），用于更新模型参数（w和b）
optimizer = torch.optim.SGD(
    model.parameters(),  # 自动获取模型中所有可训练参数（所有fc层的w和b）
    lr=0.01,  # 学习率：控制每次参数更新的幅度（0.01为经典值，过大会震荡，过小会收敛慢）
    momentum=0.5  # 动量：利用前一次梯度的方向加速收敛，减少震荡（0.5为常用值）
)


# -------------------------- 4. 训练函数：定义单轮训练逻辑（遍历所有训练样本） --------------------------
def train(epoch):
    running_loss = 0.0  # 累计当前批次的损失值（用于打印中间结果）

    # 遍历训练集DataLoader：batch_idx为批次索引，data为(输入张量, 标签张量)
    for batch_idx, data in enumerate(train_loader, 0):
        inputs, labels = data  # 解包数据：inputs（批次图像），labels（批次真实标签）

        # 步骤1：清零梯度（关键！避免上一批次梯度累积到当前批次）
        optimizer.zero_grad()
        # 步骤2：前向传播：将输入传入模型，得到预测输出（logits）
        outputs = model(inputs)
        # 步骤3：计算损失：对比预测输出与真实标签的差距
        loss = criterion(outputs, labels)
        # 步骤4：反向传播：自动计算损失对所有参数的梯度（链式法则）
        loss.backward()
        # 步骤5：参数更新：根据梯度和学习率，更新模型权重和偏置
        optimizer.step()

        running_loss += loss.item()  # 累加当前批次的损失值（.item()将张量转为Python数值，避免内存占用）

        # 每100个批次打印一次中间结果（避免输出过多，便于观察训练趋势）
        if batch_idx % 100 == 0:
            # 打印格式：[当前轮次, 批次索引] 平均损失（100个批次的平均）
            print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 100))
            running_loss = 0.0  # 重置累计损失，准备下一个100批次

    # 返回当前轮次的平均损失（总损失/批次数），用于后续绘图
    return running_loss / len(train_loader)


# -------------------------- 5. 测试函数：定义模型评估逻辑（计算测试集准确率） --------------------------
def test():
    correct = 0  # 正确预测的样本数
    total = 0  # 总测试样本数

    # 关闭梯度计算（测试阶段无需反向传播，节省内存和计算资源）
    with torch.no_grad():
        # 遍历测试集DataLoader（无梯度计算）
        for data in test_loader:
            images, labels = data  # 解包测试数据
            outputs = model(images)  # 前向传播：得到预测logits

            # 提取预测结果：torch.max(outputs.data, dim=1)返回（最大值，索引）
            # dim=1：按行取最大值（每个样本的10个类别中取概率最大的）
            # 索引即预测的类别（0-9），最大值为该类别的logits值
            _, predicted = torch.max(outputs.data, dim=1)

            total += labels.size(0)  # 累加总样本数（labels.size(0)为当前批次样本数）
            # 累加正确预测数：(predicted == labels)生成布尔张量，sum().item()统计True的数量
            correct += (predicted == labels).sum().item()

    # 打印测试集准确率（正确数/总数 * 100，保留整数）
    print('Accuracy of the network on the test images: %d %%' % (100 * correct / total))
    # 返回准确率（0-1之间），用于后续绘图
    return correct / total


# -------------------------- 6. 主程序：执行训练、测试和可视化 --------------------------
if __name__ == '__main__':
    # 初始化列表：存储训练轮次、训练损失、测试准确率（用于绘图）
    epoch_list = []  # 训练轮次（x轴数据）
    loss_list = []  # 每轮训练平均损失（y轴1数据）
    predict_list = []  # 每轮测试准确率（y轴2数据）

    # 训练循环：迭代50轮（epoch），逐步优化模型参数
    for epoch in range(15):
        # 1. 执行训练：返回当前轮平均损失
        loss_num = train(epoch)
        # 2. 记录轮次和损失
        epoch_list.append(epoch)
        loss_list.append(loss_num)
        # 3. 执行测试：返回当前轮准确率
        predict_num = test()
        # 4. 记录准确率
        predict_list.append(predict_num)

    # -------------------------- 7. 双指标可视化：绘制损失-轮次、准确率-轮次曲线 --------------------------
    # 创建画布和主轴（ax1）：用于绘制训练损失（左侧y轴）
    fig, ax1 = plt.subplots(figsize=(10, 6))  # figsize=(宽,高)：设置画布尺寸，避免标签拥挤

    # 绘制训练损失曲线（主轴ax1）
    color = 'tab:red'  # 损失曲线颜色（红色）
    ax1.set_xlabel('Epoch (训练轮次)', fontsize=12)  # x轴标签（字体大小12）
    ax1.set_ylabel('Training Loss (训练损失)', color=color, fontsize=12)  # 左侧y轴标签（红色）
    # 绘制曲线：x=epoch_list，y=loss_list，红色实线、线宽2、圆形标记（便于定位数据点）
    ax1.plot(epoch_list, loss_list, color=color, linewidth=2, marker='o', markersize=5, label='Loss')
    ax1.tick_params(axis='y', labelcolor=color)  # 左侧y轴刻度颜色设为红色（对应曲线）
    ax1.set_ylim(0, max(loss_list) * 1.1)  # 调整左侧y轴范围：从0到最大损失的1.1倍，避免曲线贴边

    # 创建副轴（ax2）：与ax1共享x轴，用于绘制测试准确率（右侧y轴）
    ax2 = ax1.twinx()  # twinx()：共享x轴，生成独立的右侧y轴
    color = 'tab:blue'  # 准确率曲线颜色（蓝色）
    ax2.set_ylabel('Test Accuracy (测试准确率)', color=color, fontsize=12)  # 右侧y轴标签（蓝色）
    # 绘制准确率曲线：蓝色实线、线宽2、方形标记
    ax2.plot(epoch_list, predict_list, color=color, linewidth=2, marker='s', markersize=5, label='Accuracy')
    ax2.tick_params(axis='y', labelcolor=color)  # 右侧y轴刻度颜色设为蓝色
    ax2.set_ylim(0, 1.0)  # 准确率范围固定为0-1（符合概率逻辑）

    # 添加图表标题和网格
    plt.title('MNIST Training Loss vs Test Accuracy', fontsize=14, pad=20)  # 标题（字体14，上下间距20）
    ax1.grid(True, linestyle='--', alpha=0.7)  # 显示网格线：虚线、透明度0.7（便于读取数值）

    # 合并图例：由于曲线分别在ax1和ax2上，需手动收集两条曲线的图例
    lines1, labels1 = ax1.get_legend_handles_labels()  # 获取ax1的曲线和标签
    lines2, labels2 = ax2.get_legend_handles_labels()  # 获取ax2的曲线和标签
    ax1.legend(lines1 + lines2, labels1 + labels2, loc='best', fontsize=10)  # 合并后显示在最佳位置

    # 调整布局：自动适配所有元素，避免标签被截断（如y轴标签超出画布）
    plt.tight_layout()

    # 保存图像（可选）：保存为高清PNG格式（dpi=300），bbox_inches='tight'避免标签截断
    plt.savefig('mnist_loss_accuracy.png', dpi=300, bbox_inches='tight')

    # 显示图像：弹出窗口展示绘制的曲线
    plt.show()